
import json
from datetime import datetime, timedelta
import numpy as np
import os
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import tools


#对overlap_ratios，绘制total_overlap_area_valid_sky随time的变化曲线
def draw_overlap_ratios(gwacSummaryFile, savePath):
    with open(gwacSummaryFile, "r") as f:
        gwacSummary = json.load(f)
    
    times = [datetime.fromisoformat(d["time"]) for d in gwacSummary['overlap_ratios']]
    # overlap_areas1 = [d["total_overlap_area_valid_sky"] for d in gwacSummary['overlap_ratios']]
    overlap_areas2 = [d["total_overlap_area_10mount_sky"] for d in gwacSummary['overlap_ratios']]
    localZenithOverlap30 = [d["localZenithOverlap30"] for d in gwacSummary['overlap_ratios']]
    mount_active_ratio = [d["validSkyNum"]/10 for d in gwacSummary['overlap_ratios']] #workingMountNum
    svom_workplans = [d for d in gwacSummary['svom_workplans']]

    date_str = times[0].strftime('%Y-%m-%d')
    max_time = max(times) + timedelta(hours=0.5)
    min_time = min(times) - timedelta(hours=0.5)

    # 创建两个子图
    fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(12, 3), sharex=True)
    axes = axes.flatten()
    
    yticks = [0.25, 0.5, 0.75]  # 设定Y轴刻度
    for i, ax in enumerate(axes):
        # ax.set_yticks(yticks)  # 设置Y轴刻度
        ax.set_ylim(0, 1.1)  # 固定Y轴范围
        ax.set_xlim(min_time, max_time)  # 固定Y轴范围
        if i<2:
            for y in yticks:
                ax.axhline(y=y, color='gray', linestyle='--', alpha=0.6)
        ax.grid()
    
    # axes[0].plot(times, overlap_areas1, linestyle='-', color='b', label='Effective Observation Sky of GWAC In SVOM')
    axes[0].plot(times, overlap_areas2, linestyle='-', color='r')
    axes[0].set_ylabel('Overlap Ratio')
    axes[0].set_title('10 Mount of GWAC In SVOM')
    # axes[0].legend()
    
    # axes[1].plot(times, overlap_areas2, linestyle='-', color='r', label='10 Mount In SVOM')
    axes[1].plot(times, localZenithOverlap30, linestyle='-', color='g')
    axes[1].set_ylabel('Overlap Ratio')
    axes[1].set_title('Horizon of 30 degrees in SVOM')
    
    axes[2].plot(times, mount_active_ratio, linestyle='-', color='g')
    axes[2].set_ylabel('Mount Active')
    axes[2].set_title('Mount Active Ratio')
    
    categories = ['GP', 'TOO-NOM', 'TOO-EX'] #'OTHER',
    category_indices = {cat: i for i, cat in enumerate(categories)}
    for item in svom_workplans:
        start = datetime.fromisoformat(item['start'])
        end = datetime.fromisoformat(item['end'])
        category = item['category']
        y_pos = category_indices.get(category, -1)
        if y_pos != -1:
            axes[3].plot([start, end], [y_pos, y_pos], linewidth=4, label=category if y_pos == 0 else "")
            
    axes[3].set_yticks(range(len(categories)))
    axes[3].set_yticklabels(categories)
    axes[3].set_ylim(-0.5, len(categories) - 0.5)  # 上下留点边界
    axes[3].grid(True)
    
    axes[3].xaxis.set_major_locator(mdates.HourLocator(interval=1))
    axes[3].xaxis.set_major_formatter(mdates.DateFormatter('%H'))  # X 轴时间格式：小时:分钟
    axes[3].set_title('SVOM Workplan')
    # axes[3].set_xlabel(f'UTC Hour of {date_str}')  # X 轴标签添加日期
    
    fig.suptitle(f'UTC Hour of {date_str}', fontsize=12)  # 设置整个图像的标题
    # 调整日期显示
    # plt.xticks(rotation=45)
    plt.tight_layout()
    # plt.show()
    plt.savefig(savePath, dpi=300)
    plt.close('all')
    print(f'save image to {savePath}')

def draw_overlap_ratios_byObsType(gwacSummaryFile):
    
    with open(gwacSummaryFile, "r") as f:
        gwacSummary = json.load(f)
    
    times = [datetime.fromisoformat(d["time"]) for d in gwacSummary['overlap_ratios']]
    overlap_areas1 = [d["total_overlap_area_valid_sky"] for d in gwacSummary['overlap_ratios']]
    overlap_areas2 = [d["total_overlap_area_10mount_sky"] for d in gwacSummary['overlap_ratios']]
    mount_active_ratio = [d["validSkyNum"]/10 for d in gwacSummary['overlap_ratios']]  
    svom_obs_types = [d["svom_obs_type"] for d in gwacSummary['overlap_ratios']]  

    date_str = times[0].strftime('%Y-%m-%d')

    # 手动指定颜色映射
    type_color_map = {'TOO-EX': 'red', 'TOO-NOM': 'green', 'OTHER': 'blue', 'GP': 'purple'}

    # 创建子图
    fig, axes = plt.subplots(nrows=3, ncols=1, figsize=(8, 5), sharex=True)

    yticks = [0.5]  
    for ax in axes:
        ax.set_yticks(yticks)
        ax.set_ylim(0, 1)
        ax.axhline(y=0.5, color='gray', linestyle='--', alpha=0.6)
        ax.grid()

    # 绘制曲线
    axes[0].plot(times, overlap_areas1, linestyle='-', color='gray', alpha=0.5)  
    axes[1].plot(times, overlap_areas2, linestyle='-', color='gray', alpha=0.5)
    axes[2].plot(times, mount_active_ratio, linestyle='-', color='gray', alpha=0.5)

    # 根据类型上色
    for i, t in enumerate(svom_obs_types):
        color = type_color_map.get(t, 'gray')  # 处理未知类型
        axes[0].scatter(times[i], overlap_areas1[i], color=color, s=10)
        axes[1].scatter(times[i], overlap_areas2[i], color=color, s=10)
        axes[2].scatter(times[i], mount_active_ratio[i], color=color, s=10)

    # 设置标题和标签
    axes[0].set_ylabel('Overlap Ratio')
    axes[0].set_title('Valid Observation Sky of GWAC In SVOM')

    axes[1].set_ylabel('Overlap Ratio')
    axes[1].set_title('10 Mount In SVOM')

    axes[2].set_ylabel('Mount Active')
    axes[2].set_title('Mount Active Ratio')

    # 设置 X 轴
    axes[2].xaxis.set_major_locator(mdates.HourLocator(interval=1))  
    axes[2].xaxis.set_major_formatter(mdates.DateFormatter('%H:%M'))  
    axes[2].set_xlabel(f'Time (UTC) - {date_str}')

    # 添加图例
    legend_patches = [plt.Line2D([0], [0], marker='o', color=color, linestyle='', markersize=6, label=t) for t, color in type_color_map.items()]
    axes[0].legend(handles=legend_patches, title="svom_obs_type")

    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.show()

def plot_overlap_xinglong(overlap_areas):
    times = [datetime.fromisoformat(entry[0]) for entry in overlap_areas]
    areas = [entry[1] for entry in overlap_areas]
    
    date_str = times[0].strftime('%Y-%m-%d')

    plt.figure(figsize=(10, 5))
    plt.plot(times, areas, linestyle='-', color='b', label="overlap area")

    plt.xlabel("Time (UTC)")
    plt.ylabel("Overlap area")
    plt.title("Xiong and SVOM ")
    plt.legend()
    
    # 美化时间刻度
    plt.gca().xaxis.set_major_locator(mdates.HourLocator(interval=1))
    plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%H'))
    plt.gcf().autofmt_xdate()
    plt.gca().set_xlabel(f'UTC Hour of {date_str}')  # X 轴标签添加日期

    plt.grid(True, linestyle="--", alpha=0.6)
    plt.show()

def test():
    
    # utcTimeStr = tools.getCurDateStr()
    utcTimeStr = "2025-02-26"
    dataDir = f"public/svomSummary/data/{utcTimeStr}"
    gwacSummaryFile = f"{dataDir}/observation_results.json"
    
    if os.path.exists(gwacSummaryFile):
        savePath = f"{dataDir}/overlap_ratios.png"
        draw_overlap_ratios(gwacSummaryFile, savePath)
        # draw_overlap_ratios_byObsType(gwacSummaryFile)
    else:
        print(f"{gwacSummaryFile} not exist.")
    
if __name__ == "__main__":
    
    test()